import sqlite3
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_absolute_error
import matplotlib.pyplot as plt
from sklearn.tree import plot_tree, DecisionTreeRegressor
import streamlit as st
import numpy as np

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 读取数据
data = pd.read_csv('house_price.csv')

# 特征工程 - 对位置进行编码
le = LabelEncoder()
data['place_encoded'] = le.fit_transform(data['place'])

# 准备特征和目标变量
X = data[['area', 'age', 'place_encoded']]  # 使用所有特征
y = data['price']

# 数据标准化
transfer = StandardScaler()
X_scaled = transfer.fit_transform(X)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=42)

# 训练线性回归模型
lr = LinearRegression()
lr.fit(X_train, y_train)
y_pred_lr = lr.predict(X_test)

# 训练决策树模型
dt = DecisionTreeRegressor(random_state=42)
dt.fit(X_train, y_train)
y_pred_dt = dt.predict(X_test)

# 计算模型性能
lr_r2 = r2_score(y_test, y_pred_lr)
lr_mae = mean_absolute_error(y_test, y_pred_lr)
dt_r2 = r2_score(y_test, y_pred_dt)
dt_mae = mean_absolute_error(y_test, y_pred_dt)

# Streamlit界面
st.set_page_config(page_title="房屋价格预测系统", page_icon="🏠", layout="wide")
st.title("🏠 房屋价格预测系统")

# 预测部分
col1, col2, col3 = st.columns(3)
with col1:
    area = st.number_input('房屋面积（平方米）', 0, 500, 100, 10)

with col2:
    age = st.number_input('房龄（年）', 0, 50, 5, 1)

with col3:
    place_options = data['place'].unique().tolist()
    place = st.selectbox('位置', place_options)
    place_encoded = le.transform([place])[0]

# 准备输入数据
input_data = np.array([[area, age, place_encoded]])
input_scaled = transfer.transform(input_data)

if st.button('预测房价'):
    pred_lr = lr.predict(input_scaled)[0]
    pred_dt = dt.predict(input_scaled)[0]

    # 显示预测结果
    st.success("预测完成！")

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("线性回归预测")
        st.metric("预测价格", f"¥{pred_lr:,.2f}",
                  delta=f"{(pred_lr - data['price'].mean()) / data['price'].mean() * 100:.1f}%")
        st.write(f"R²分数: {lr_r2:.3f}")
        st.write(f"平均绝对误差: ¥{lr_mae:,.2f}")

    with col2:
        st.subheader("决策树预测")
        st.metric("预测价格", f"¥{pred_dt:,.2f}",
                  delta=f"{(pred_dt - data['price'].mean()) / data['price'].mean() * 100:.1f}%")
        st.write(f"R²分数: {dt_r2:.3f}")
        st.write(f"平均绝对误差: ¥{dt_mae:,.2f}")

# 模型可视化部分
st.markdown("---")
st.title("房地产价格预测模型可视化")

# 创建选项卡
tab1, tab2 = st.tabs(["决策树可视化", "模型性能"])

with tab1:
    st.header("决策树结构可视化")

    # 添加控制参数
    col1, col2 = st.columns(2)
    with col1:
        max_depth = st.slider("最大深度", 1, 10, 3, key="tree_depth")
    with col2:
        font_size = st.slider("字体大小", 6, 16, 10, key="font_size")

    # 重新训练决策树模型以应用新参数
    dt_viz = DecisionTreeRegressor(max_depth=max_depth, random_state=42)
    dt_viz.fit(X_train, y_train)

    # 绘制决策树
    fig, ax = plt.subplots(figsize=(16, 10))
    plot_tree(dt_viz,
              feature_names=['面积', '房龄', '位置编码'],
              filled=True,
              rounded=True,
              proportion=True,
              fontsize=font_size,
              ax=ax)
    ax.set_title(f"决策树结构 (最大深度: {max_depth})", fontsize=16, pad=20)
    plt.tight_layout()
    st.pyplot(fig)
    plt.close(fig)

    # 显示决策树性能
    y_pred_dt_viz = dt_viz.predict(X_test)
    dt_viz_r2 = r2_score(y_test, y_pred_dt_viz)
    dt_viz_mae = mean_absolute_error(y_test, y_pred_dt_viz)

    st.info(f"""
    **当前决策树性能:**
    - R²分数: {dt_viz_r2:.3f}
    - 平均绝对误差: ¥{dt_viz_mae:,.2f}
    - 树深度: {dt_viz.get_depth()}
    - 叶节点数: {dt_viz.get_n_leaves()}
    """)

with tab2:
    st.header("模型性能比较")

    # 性能对比
    performance_data = {
        '模型': ['线性回归', '决策树'],
        'R²分数': [lr_r2, dt_r2],
        '平均绝对误差': [lr_mae, dt_mae]
    }

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("R²分数比较")
        for model, score in zip(['线性回归', '决策树'], [lr_r2, dt_r2]):
            st.metric(f"{model} R²", f"{score:.3f}")

    with col2:
        st.subheader("平均绝对误差比较")
        for model, mae in zip(['线性回归', '决策树'], [lr_mae, dt_mae]):
            st.metric(f"{model} MAE", f"¥{mae:,.2f}")

    # 特征重要性（决策树）
    st.subheader("特征重要性")
    feature_importance = dt.feature_importances_
    feature_names = ['面积', '房龄', '位置编码']

    importance_df = pd.DataFrame({
        '特征': feature_names,
        '重要性': feature_importance
    }).sort_values('重要性', ascending=False)

    st.dataframe(importance_df.style.format({'重要性': '{:.3f}'}))

    # 线性回归系数
    st.subheader("线性回归系数")
    coef_df = pd.DataFrame({
        '特征': feature_names,
        '系数': lr.coef_
    })
    st.dataframe(coef_df.style.format({'系数': '{:.3f}'}))

# 数据探索部分
st.markdown("---")
st.header("数据探索")

col1, col2 = st.columns(2)

with col1:
    st.subheader("数据统计摘要")
    st.dataframe(data.describe().style.format('{:.2f}'))

with col2:
    st.subheader("原始数据")
    st.dataframe(data)

# 添加一些数据可视化
st.subheader("数据分布")
fig, axes = plt.subplots(1, 2, figsize=(12, 4))

# 价格分布
axes[0].hist(data['price'], bins=10, alpha=0.7, color='skyblue')
axes[0].set_xlabel('价格')
axes[0].set_ylabel('频次')
axes[0].set_title('价格分布')

# 面积分布
axes[1].hist(data['area'], bins=10, alpha=0.7, color='lightgreen')
axes[1].set_xlabel('面积')
axes[1].set_ylabel('频次')
axes[1].set_title('面积分布')

plt.tight_layout()
st.pyplot(fig)
plt.close(fig)